#include "stdafx.h"
#include "TCPSession.h"
#include "HttpHeader.h"
#include <fstream>

using namespace std;

#if 0//_MSC_PLATFORM_TOOLSET > 100
std::mutex CTCPSessionMgr::gPoolLock;
#else
CRITICAL_SECTION CTCPSessionMgr::gPoolLock;
#endif

std::unordered_map<std::string, CTcpSession*> CTCPSessionMgr::gSessionPool;

CTCPSessionMgr::CTCPSessionMgr()
{
}


CTCPSessionMgr::~CTCPSessionMgr()
{
}

CTcpSession * CTCPSessionMgr::AddSession(const iphdr_t * iphdr, const tcphdr_t * tcphdr)
{
#if _MSC_PLATFORM_TOOLSET > 100
	string key1 = std::to_string(iphdr->ip_dst.S_un.S_addr) + to_string(tcphdr->th_dport) +
		to_string(iphdr->ip_src.S_un.S_addr) + to_string(tcphdr->th_sport);
	string key2 = to_string(iphdr->ip_src.S_un.S_addr) + to_string(tcphdr->th_sport) +
		std::to_string(iphdr->ip_dst.S_un.S_addr) + to_string(tcphdr->th_dport);
#else 
	char key[MAX_PATH] = { 0 };
	sprintf_s(key, "%u%u%u%u", iphdr->ip_dst.S_un.S_addr, tcphdr->th_dport, iphdr->ip_src.S_un.S_addr, tcphdr->th_sport);
	string key1 = key;
	sprintf_s(key, "%u%u%u%u", iphdr->ip_src.S_un.S_addr, tcphdr->th_sport, iphdr->ip_dst.S_un.S_addr, tcphdr->th_dport);
	string key2 = key;
#endif

	CTcpSession *session = nullptr;
	{
#if 0//_MSC_PLATFORM_TOOLSET > 100
		lock_guard<mutex> lk(gPoolLock);
#else
		EnterCriticalSection(&gPoolLock);
#endif

		unordered_map<string, CTcpSession*>::iterator it1 = gSessionPool.find(key1);
		if (it1 != gSessionPool.end()) {
			//Error condition
			LeaveCriticalSection(&gPoolLock);
			return nullptr;
		}
		unordered_map<string, CTcpSession*>::iterator it2 = gSessionPool.find(key2);
		if (it2 != gSessionPool.end()) {
			LeaveCriticalSection(&gPoolLock);
			return it2->second;
		}
		session = new CTcpSession(iphdr, tcphdr);
		if (!session){
			LeaveCriticalSection(&gPoolLock);
			return nullptr;
		}

		gSessionPool.insert(make_pair(key1, session));
	}
	LeaveCriticalSection(&gPoolLock);
	return session;
}

CTcpSession * CTCPSessionMgr::GetSession(const iphdr_t * iphdr, const tcphdr_t * tcphdr)
{
#if _MSC_PLATFORM_TOOLSET > 100
	string key1 = std::to_string(iphdr->ip_dst.S_un.S_addr) + to_string(tcphdr->th_dport) +
		to_string(iphdr->ip_src.S_un.S_addr) + to_string(tcphdr->th_sport);
	string key2 = to_string(iphdr->ip_src.S_un.S_addr) + to_string(tcphdr->th_sport) +
		std::to_string(iphdr->ip_dst.S_un.S_addr) + to_string(tcphdr->th_dport);
#else 
	char key[MAX_PATH] = { 0 };
	sprintf_s(key, "%u%u%u%u", iphdr->ip_dst.S_un.S_addr, tcphdr->th_dport, iphdr->ip_src.S_un.S_addr, tcphdr->th_sport);
	string key1 = key;
	sprintf_s(key, "%u%u%u%u", iphdr->ip_src.S_un.S_addr, tcphdr->th_sport, iphdr->ip_dst.S_un.S_addr, tcphdr->th_dport);
	string key2 = key;
#endif

	{
#if 0//_MSC_PLATFORM_TOOLSET > 100
		lock_guard<mutex> lk(gPoolLock);
#else
		EnterCriticalSection(&gPoolLock);
#endif

		unordered_map<string, CTcpSession*>::iterator it1 = gSessionPool.find(key1);
		if (it1 != gSessionPool.end()) {
			LeaveCriticalSection(&gPoolLock);
			return it1->second;
		}

		unordered_map<string, CTcpSession*>::iterator it2 = gSessionPool.find(key2);
		if (it2 != gSessionPool.end()) {
			LeaveCriticalSection(&gPoolLock);
			return it2->second;
		}
	}
	LeaveCriticalSection(&gPoolLock);
	return nullptr;
}

bool CTCPSessionMgr::CloseSession(const iphdr_t * iphdr, const tcphdr_t * tcphdr)
{
#if _MSC_PLATFORM_TOOLSET > 100
	string key1 = std::to_string(iphdr->ip_dst.S_un.S_addr) + to_string(tcphdr->th_dport) +
		to_string(iphdr->ip_src.S_un.S_addr) + to_string(tcphdr->th_sport);
	string key2 = to_string(iphdr->ip_src.S_un.S_addr) + to_string(tcphdr->th_sport) +
		std::to_string(iphdr->ip_dst.S_un.S_addr) + to_string(tcphdr->th_dport);
#else 
	char key[MAX_PATH] = { 0 };
	sprintf_s(key, "%u%u%u%u", iphdr->ip_dst.S_un.S_addr, tcphdr->th_dport, iphdr->ip_src.S_un.S_addr, tcphdr->th_sport);
	string key1 = key;
	sprintf_s(key, "%u%u%u%u", iphdr->ip_src.S_un.S_addr, tcphdr->th_sport, iphdr->ip_dst.S_un.S_addr, tcphdr->th_dport);
	string key2 = key;
#endif

	CTcpSession *session = nullptr;
	{
#if 0//_MSC_PLATFORM_TOOLSET > 100
		lock_guard<mutex> lk(gPoolLock);
#else
		EnterCriticalSection(&gPoolLock);
#endif

		unordered_map<string, CTcpSession*>::iterator it1 = gSessionPool.find(key1);
		if (it1 == gSessionPool.end()) {
			unordered_map<string, CTcpSession*>::iterator it2 = gSessionPool.find(key2);
			if (it2 == gSessionPool.end()) {
				LeaveCriticalSection(&gPoolLock);
				return false;
			}
			else {
				session = it2->second;
				gSessionPool.erase(it2);
			}
		}
		else {
			session = it1->second;
			gSessionPool.erase(it1);
		}

		delete session;
	}
	LeaveCriticalSection(&gPoolLock);
	return true;
}

bool CTCPSessionMgr::Initial()
{
#if 1
	InitializeCriticalSection(&gPoolLock);
#endif
	return true;
}


CTcpSession::CTcpSession(const iphdr_t*iphdr, const tcphdr_t *tcphdr) :mIpHdr(*iphdr), mTcpHdr(*tcphdr), mUpLinkConnect(NULL),
mDownLinkConnect(NULL), mSessionState(CnState_NONE), mBuffer(NULL), mBufflen(0), mPos(0), mChunked(FALSE), mChunkOffset(0)
{
	//printf("New Connect\n");
	char name[MAX_PATH] = { 0 };
	sprintf_s(name, "./dump/%u.%u.%u.%u-%u_%u.%u.%u.%u-%u.txt", iphdr->ip_dst.S_un.S_un_b.s_b1, iphdr->ip_dst.S_un.S_un_b.s_b2, iphdr->ip_dst.S_un.S_un_b.s_b3, iphdr->ip_dst.S_un.S_un_b.s_b4, ntohs(tcphdr->th_dport),
		iphdr->ip_src.S_un.S_un_b.s_b1, iphdr->ip_src.S_un.S_un_b.s_b2, iphdr->ip_src.S_un.S_un_b.s_b3, iphdr->ip_src.S_un.S_un_b.s_b4, ntohs(tcphdr->th_sport));
	mUpLinkConnect = new CTcpConnect(name, TCPBUFSIZE);
	sprintf_s(name, "./dump/%u.%u.%u.%u-%u_%u.%u.%u.%u-%u.txt", iphdr->ip_src.S_un.S_un_b.s_b1, iphdr->ip_src.S_un.S_un_b.s_b2, iphdr->ip_src.S_un.S_un_b.s_b3, iphdr->ip_src.S_un.S_un_b.s_b4, ntohs(tcphdr->th_sport),
		iphdr->ip_dst.S_un.S_un_b.s_b1, iphdr->ip_dst.S_un.S_un_b.s_b2, iphdr->ip_dst.S_un.S_un_b.s_b3, iphdr->ip_dst.S_un.S_un_b.s_b4, ntohs(tcphdr->th_dport));
	mDownLinkConnect = new CTcpConnect(name, TCPBUFSIZE);
}


CTcpSession::~CTcpSession()
{
	if (mUpLinkConnect)
		delete mUpLinkConnect;
	if (mDownLinkConnect)
		delete mDownLinkConnect;

	if (mBuffer) {
		delete[]mBuffer;
		mBuffer = NULL;
	}
}

void CTcpSession::Dump2File(string& filename, char*buf, INT32 size, ios_base::openmode mode)
{
	if (filename.length() == 0)
		return;

	DEBUG_PRINT("WriteFile:%s\n", filename.c_str());
	fstream file;
	file.open(filename.c_str(), mode);
	if (file.is_open()) {
		file.write(mBuffer, mPos);
	}
	else {
		printf("Create file failed!\n");
	}
}

INT32 CTcpSession::CopyPayload2Buffer(char*buffer, INT32 bufSize, INT32 &bufPos, const uchar*payload, INT32 pldLen, INT32 &offset, BOOL chunked)
{
	//assume this packet is ordered
	if (!chunked) {
		if (bufSize - bufPos < pldLen) {
			printf("ErrorCondition: bufSize[%d] - bufPos[%d] < pldLen[%d]\n", bufSize, bufPos, pldLen);
			return -1;//Error argument
		}
		memcpy_s(&buffer[bufPos], bufSize - bufPos, payload, pldLen);
		bufPos += pldLen;
		return 0;
	}

	INT32 copylen = 0;
	INT32 pldoffset = 0;
	size_t pos = string::npos;
	do {
		if (offset <= 0) {
			string chunkhead = reinterpret_cast<char*>(const_cast<uchar*>(&payload[pldoffset]));
			if ((pos = chunkhead.find("\r\n")) == string::npos) {
				printf("ErrorCondition: error ChunkHeader\n");
				return -1;//Error argument
			}
			chunkhead = chunkhead.substr(0, pos);
			offset = strtol(chunkhead.c_str(), NULL, 16);
			if (offset < 0) {
				printf("ErrorCondition: error ChunkHeader len\n");
				return -1;//Error argument
			}
			pldoffset += static_cast<INT32>(pos) + 2;
			pldLen -= static_cast<INT32>(pos) + 2;
		}

		if (offset == 0) {
			return 1;//completed
		}

		copylen = pldLen > offset ? offset : pldLen;
		INT32 bufleft = bufSize - bufPos;

		copylen = (bufleft <= copylen) ? bufleft : copylen;
		memcpy_s(&buffer[bufPos], bufleft, &payload[pldoffset], copylen);
		bufPos += copylen;
		pldoffset += copylen;
		offset -= copylen;
		pldLen -= copylen;
		if (offset == 0) {
			pldoffset += 2;
			pldLen -= 2;
		}
		if (bufSize - bufPos <= 0) {
			Dump2File(mRequestContent, mBuffer, mPos, ios_base::out | ios_base::binary | ios_base::ate | ios_base::app);
			if (mBuffer) {
				mPos = 0;
			}
		}
	} while (pldLen);
	//memcpy_s(mBuffer, mBufflen, &payload[pos + 4 + pos + 2], chuncklen);
	//mPos += chuncklen;
	return 0;
}

INT32 CTcpSession::HandlePacket(CETHFram*pkt, const iphdr_t*iphdr, const tcphdr_t *tcphdr, const uchar * payload, const uint32_t payloadlen)
{
	CTcpConnect *CurConnect = NULL;
	LinkMode mode = CnDir_MAX;
	if (iphdr->ip_dst.S_un.S_addr == mIpHdr.ip_dst.S_un.S_addr && iphdr->ip_src.S_un.S_addr == mIpHdr.ip_src.S_un.S_addr&&
		tcphdr->th_dport == mTcpHdr.th_dport && tcphdr->th_sport == mTcpHdr.th_sport) {//Same with SYN flag packet
		mode = CnDir_Up;
		CurConnect = mUpLinkConnect;
	}
	else {
		mode = CnDir_Down;
		CurConnect = mDownLinkConnect;
	}
	if (!CurConnect)//delete this session just because error session
		return -1;


	switch (mSessionState)
	{
	case CnState_NONE:
		if (mode == CnDir_Down) {//Ignore packet form server when NONE state
			return 0;
		}
		mSessionState = CnState_HandShark;
	case CnState_HandShark:
		if (mode == CnDir_Up) {
			//Handle Get Request
			// filter URL

			CHttpRequest request;
			if (request.PreProcess(payload, payloadlen) == FALSE) {
				mSessionState = CnState_NONE;
				return 0;
			}
			if (request.mRequestLine.strMethod == "GET") {
				DEBUG_PRINT("\tMethod:%s\n\tContent:%s\n\tPostFix:%s\n", request.mRequestLine.strMethod.c_str(),
					request.mRequestLine.RequestContent().c_str(), request.mRequestLine.RequestPostfix().c_str());
				if (request.mRequestLine.RequestPostfix().length() == 0 ||
					(request.mRequestLine.RequestPostfix() != "jpg")) {
					//&& request.mRequestLine.RequestPostfix() != "png"&& request.mRequestLine.RequestPostfix() != "png"
					mSessionState = CnState_NONE;
					break;
				}

				mRequestContent = "./Download/" + request.mRequestLine.RequestContent();
			}
			else
				mSessionState = CnState_NONE;

			break;
		}
		else {
			//Handle Http Response
			// Check result, content length, decoding
			CHttpResponse response;
			if (response.PreProcess(payload, payloadlen) == FALSE) {
				DEBUG_PRINT("ip.addr==%s && ip.addr==%s && tcp.port==%u && tcp.port==%u\n", inet_ntoa(iphdr->ip_src), inet_ntoa(iphdr->ip_dst),
					ntohs(tcphdr->th_sport), ntohs(tcphdr->th_dport));
				mSessionState = CnState_NONE;
				return 0;
			}

			DEBUG_PRINT("Response:\n\t Status: %s\n\t Reason: %s\n", response.mResponseLine.StatusCode.c_str(), response.mResponseLine.reason.c_str());
			if (response.mResponseLine.StatusCode != "200") {
				mSessionState = CnState_NONE;
				return 0;
			}

			if (response.mContentLength <= 8196) {
				if (response.mChunked == false) {
					mSessionState = CnState_NONE;
					return 0;
				}
				else {
					mChunked = TRUE;
				}
			}
			else {
				mChunked = FALSE;
			}

			string data(reinterpret_cast<const char*>(payload));
			size_t pos = data.find("\r\n\r\n");
			if (pos == string::npos) {
				mSessionState = CnState_NONE;
				break;
			}

			int datalen = payloadlen - static_cast<INT32>(pos) - 4;
			if (response.mContentLength == datalen) {
				//printf("ContentLength:%u  datalen:%u\n", response.mContentLength, datalen);
				//Complete 
				Dump2File(mRequestContent, reinterpret_cast<char*>(const_cast<uchar*>(&payload[pos + 4])),
					datalen, ios_base::trunc | ios_base::out | ios_base::binary);

				mSessionState = CnState_HandShark;
			}
			else {
				if (mBuffer) {
					printf("Maybe lost packet\n");
					delete[]mBuffer;
				}
				if (!mChunked) {
					mBuffer = new char[response.mContentLength];
					mBufflen = response.mContentLength;
					mPos = 0;
				}
				else {
					mBuffer = new char[DEFAULTBUFFERSIZE];
					mBufflen = DEFAULTBUFFERSIZE;
					mPos = 0;
					mChunkOffset = 0;
				}
				int ret = CopyPayload2Buffer(mBuffer, mBufflen, mPos, &payload[pos + 4], datalen, mChunkOffset, mChunked);
				if (ret != 0) {
					printf("ErrorCondition: copypayload of first http response!!!\n");
				}

				mExpectSeq = pkt->GetPktSeq() + pkt->GetTCPPayloadLen();
				DEBUG_PRINT("First Packet: payload:%u expectSeq:%u\n", mPos, mExpectSeq);

				mSessionState = CnState_WaitPacket;
			}

			break;
		}
	case CnState_WaitPacket:
		if (mode == CnDir_Down) {
			//handle packet
			//Store all content for now
			if (pkt->GetPktSeq() != mExpectSeq) {
				DEBUG_PRINT("Unordered Packet: Seq:%u ExpectSeq:%u\n", pkt->GetPktSeq(), mExpectSeq);
				CurConnect->AddPacket(pkt);
				return 1;
			}

			int ret = CopyPayload2Buffer(mBuffer, mBufflen, mPos, pkt->GetTCPPayload(), pkt->GetTCPPayloadLen(), mChunkOffset, mChunked);
			if (ret == -1) {
				goto ErrorConn;
			}
			else if (ret == 1) {
				goto FinishConn;
			}
			mExpectSeq = pkt->GetPktSeq() + pkt->GetTCPPayloadLen();
			DEBUG_PRINT("Ordered Packet: Seq:%u ExpectSeq:%u buffer:%u\n", pkt->GetPktSeq(), mExpectSeq, mPos);

			CETHFram*tmp = NULL;
			while ((tmp = CurConnect->GetExpectPacket(mExpectSeq)) != NULL) {
				ret = CopyPayload2Buffer(mBuffer, mBufflen, mPos, tmp->GetTCPPayload(), tmp->GetTCPPayloadLen(), mChunkOffset, mChunked);
				if (ret == -1) {
					goto ErrorConn;
				}
				else if (ret == 1) {
					goto FinishConn;
				}

				mExpectSeq = tmp->GetPktSeq() + tmp->GetTCPPayloadLen();
				DEBUG_PRINT("GetExpect Packet: Seq:%u ExpectSeq:%u buffer:%u\n", tmp->GetPktSeq(), mExpectSeq, mPos);

				delete tmp;
			}
			if (mChunked == false && mPos == mBufflen) {
				goto FinishConn;
			}
		}
		//Ignore the other case packet
		break;
	default:
		break;
	}
	return 0;

ErrorConn:
	printf("ErrorCondition: copypayload of payload\n");
	if (mBuffer) {
		delete[]mBuffer;
		mBuffer = NULL;
		mBufflen = 0;
		mPos = 0;
		mChunkOffset = 0;
	}

	mSessionState = CnState_NONE;
	return 0;

FinishConn:
	Dump2File(mRequestContent, mBuffer, mPos, ios_base::out | ios_base::binary | ios_base::ate | ios_base::app);

	if (mBuffer) {
		delete[]mBuffer;
		mBuffer = NULL;
		mBufflen = 0;
		mPos = 0;
	}
	CurConnect->Clear();
	mRequestContent = "";
	mSessionState = CnState_HandShark;
	return 0;
}

CTcpConnect::CTcpConnect(string filename, int32_t totalsize)// :mTotalSize(0)
{
}

CTcpConnect::~CTcpConnect()
{
	Clear();
}

BOOL CTcpConnect::AddPacket(CETHFram * pkt)
{
	if (!pkt) {
		return FALSE;
	}
	ULONG seq = pkt->GetPktSeq();
	BOOL flag = false;
	for (list<CETHFram*>::iterator it = mPktQueue.begin(); it != mPktQueue.end(); it++) {
		if ((*it)->GetPktSeq() > seq) {
			mPktQueue.insert(it, pkt);
			flag = true;
			break;
		}
		else if ((*it)->GetPktSeq() == seq) {
			return FALSE;
		}
	}

	if (flag == false) {
		mPktQueue.push_back(pkt);
	}

	return true;
}

CETHFram * CTcpConnect::GetExpectPacket(ULONG seq)
{
	CETHFram*pkt = NULL;
	for (list<CETHFram*>::iterator it = mPktQueue.begin(); it != mPktQueue.end(); it++) {
		if ((*it)->GetPktSeq() == seq) {
			pkt = *it;
			mPktQueue.erase(it);
			break;
		}
		else if ((*it)->GetPktSeq() > seq) {
			break;
		}
		else {
			DEBUG_PRINT("ErrorCondition: pkt sequence error\n");
		}
	}

	return pkt;
}

void CTcpConnect::Clear()
{
	while (mPktQueue.size()) {
		CETHFram*pkt = mPktQueue.front();
		mPktQueue.pop_front();
		delete pkt;
	}
}
